package cn.actoncode.boot.module.ai.service.chatmodel.impl;

import cn.actoncode.boot.framework.zhipu.core.bean.*;
import cn.actoncode.boot.framework.zhipu.core.client.ChatGLMApiClient;
import cn.actoncode.boot.framework.zhipu.core.enums.ChatGLMFinishReason;
import cn.actoncode.boot.framework.zhipu.core.enums.ChatGLMRoleEnum;
import cn.actoncode.boot.module.ai.service.chatmodel.dto.ChatHistoryDTO;
import cn.actoncode.boot.module.ai.service.chatmodel.dto.ChatRespDTO;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSON;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;

import static cn.actoncode.boot.framework.common.exception.enums.GlobalErrorCodeConstants.*;
import static cn.actoncode.boot.framework.common.exception.util.ServiceExceptionUtil.exception;
import static cn.actoncode.boot.framework.zhipu.core.client.ChatGLMApiClient.QPS_ERROR_CODE;
import static cn.actoncode.boot.framework.zhipu.core.client.ChatGLMApiClient.SEC_CHECK_ERROR_CODE;

@Slf4j
public abstract class AbstractChatGLM3ServiceImpl extends AbstractChatServiceImpl {

    @Resource
    protected ChatGLMApiClient chatGLMApiClient;

    @Override
    protected ChatRespDTO doChat(String prompt, List<ChatHistoryDTO> historyList, String system,
                                 Boolean stream, String userId, Long sk) {
        // 构建请求消息体
        List<ChatGLMV4Request.Message> messages = new ArrayList<>();
        if (StrUtil.isNotEmpty(system)) {
            ChatGLMV4Request.Message message = new ChatGLMV4Request.Message();
            message.setRole(ChatGLMRoleEnum.system.name());
            message.setContent(system);
            messages.add(message);
        }
        if (CollUtil.isNotEmpty(historyList)) {
            historyList.forEach(history -> {
                ChatGLMV4Request.Message message = new ChatGLMV4Request.Message();
                message.setRole(ChatGLMRoleEnum.user.name());
                message.setContent(history.getQuery());
                ChatGLMV4Request.Message message1 = new ChatGLMV4Request.Message();
                message1.setRole(ChatGLMRoleEnum.assistant.name());
                message1.setContent(history.getContent());
                messages.add(message);
                messages.add(message1);
            });
        }
        // 构建请求体
        ChatGLMV4Request chatGLMV4Request = chatGLMApiClient.buildRequestV4(prompt, stream, messages);
        if (stream) {
            // 获取流
            InputStream inputStream = this.streamChat(chatGLMV4Request);
            // 发送消息
            ChatGLMSocketResponse socketResp = sendMessage(sk, inputStream);
            return ChatRespDTO.builder().prompt(prompt).text(socketResp.getText())
                    .requestBody(JSON.toJSONString(chatGLMV4Request))
                    .promptTokens(socketResp.getMeta().getUsage().getPromptTokens())
                    .completionTokens(socketResp.getMeta().getUsage().getCompletionTokens())
                    .totalTokens(socketResp.getMeta().getUsage().getTotalTokens())
                    .isMaxLen(socketResp.getMeta().getUsage().getTotalTokens() > 10000 || StrUtil.isEmpty(socketResp.getText())).build();
        } else {
            ChatGLMV4Response chatGLMResponse = this.chat(chatGLMV4Request);
            StringBuilder sb = new StringBuilder();
            // 循环拼接结果
            chatGLMResponse.getChoices().forEach(choice -> {
                String content = choice.getMessage().getContent();
                content = content.startsWith("\"") ? content.substring(1) : content;
                content = content.endsWith("\"") ? content.substring(0, content.length() - 1) : content;
                content = content.replaceAll("\\\\n", "\n");
                sb.append(content);
            });
            String text = sb.toString();
            return ChatRespDTO.builder().prompt(prompt).text(text)
                    .requestBody(JSON.toJSONString(chatGLMV4Request))
                    .promptTokens(chatGLMResponse.getUsage().getPromptTokens())
                    .completionTokens(chatGLMResponse.getUsage().getCompletionTokens())
                    .totalTokens(chatGLMResponse.getUsage().getTotalTokens())
                    .isMaxLen(chatGLMResponse.getUsage().getTotalTokens() > 10000 || StrUtil.isEmpty(text)).build();
        }
    }

    @Override
    protected ChatGLMSocketResponse sendMessage(Long sk, InputStream inputStream) {
        BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream));
        String line;
        StringBuilder text = new StringBuilder();
        ChatGLMSocketResponse.ChatGLMMeta chatGLMMeta = null;
        try {
            while ((line = reader.readLine()) != null) {

                if (StrUtil.isEmpty(line)){
                    continue;
                }
                System.out.println(line);
                if (line.contains("data: [DONE]")) {
                    super.sendMessage(sk, "WRITE_FINISH");
                } else if (line.startsWith("{\"error\"")) {
                    ChatGLMStreamError chatGLMStreamError = JSON.parseObject(line, ChatGLMStreamError.class);
                    ChatGLMStreamError.Error error = chatGLMStreamError.getError();
                    int code = error.getCode();
                    String msg = error.getMessage();
                    log.info("NLP_ERROR：{}", msg);
                    if (ArrayUtil.contains(SEC_CHECK_ERROR_CODE, code)) {
                        throw exception(NLP_SEC_CHECK_ERROR);
                    } else if (ArrayUtil.contains(QPS_ERROR_CODE, code)) {
                        throw exception(NLP_QPS_ERROR);
                    } else {
                        throw exception(NLP_ERROR, "当前功能使用人数过多，请稍后再试");
                    }
                }  else if (line.startsWith("data:")) {
                    String data = line.substring(5);
                    ChatGLMStreamData chatGLMStreamData = JSON.parseObject(data, ChatGLMStreamData.class);
                    if (chatGLMStreamData.getUsage() != null) {
                        chatGLMMeta = new ChatGLMSocketResponse.ChatGLMMeta();
                        chatGLMMeta.setRequestId(chatGLMStreamData.getId());

                        ChatGLMUsage usage = new ChatGLMUsage();
                        usage.setPromptTokens(chatGLMStreamData.getUsage().getPromptTokens());
                        usage.setCompletionTokens(chatGLMStreamData.getUsage().getCompletionTokens());
                        usage.setTotalTokens(chatGLMStreamData.getUsage().getTotalTokens());
                        chatGLMMeta.setUsage(usage);
                    }
                    for (ChatGLMStreamData.Choices choice : chatGLMStreamData.getChoices()) {
                        if (null != choice.getFinishReason() && choice.getFinishReason().equals(ChatGLMFinishReason.sensitive.name())) {
                            throw exception(NLP_SEC_CHECK_ERROR);
                        }
                        choice.getDelta().forEach(delta -> {
                            super.sendMessage(sk, delta.getContent());
                            text.append(delta.getContent());
                        });
                    }
                }

            }

            reader.close();
            inputStream.close();
        } catch (IOException e) {
            log.info(e.getMessage());
            throw exception(NLP_ERROR, "当前功能使用人数过多，请稍后再试");
        }
        return ChatGLMSocketResponse.builder().text(text.toString()).meta(chatGLMMeta).build();
    }

    protected abstract InputStream streamChat(ChatGLMV4Request chatGLMRequest);

    protected abstract ChatGLMV4Response chat(ChatGLMV4Request chatGLMRequest);
}
